from transformers import AutoTokenizer, AutoModelForSeq2SeqLM


class Translator:
    def __init__(self):
        self.model_name = "liam168/trans-opus-mt-zh-en"
        self.tokenizer = AutoTokenizer.from_pretrained(self.model_name, local_files_only=True)
        self.model = AutoModelForSeq2SeqLM.from_pretrained(self.model_name, local_files_only=True)

    def translate_cn_to_en(self, chinese_sentences: list) -> list:
        if isinstance(chinese_sentences, str):
            chinese_sentences = [chinese_sentences]
        inputs = self.tokenizer(chinese_sentences, return_tensors="pt", padding=True, truncation=True)
        outputs = self.model.generate(**inputs)
        english_translation = [self.tokenizer.decode(o, skip_special_tokens=True) for o in outputs]
        return english_translation


if __name__ == '__main__':
    translator = Translator()
    chinese_sentences = [
        "你好，世界。",
        "草泥马",
        "张三, 李四, 王五, 赵六, 诸葛亮, 轩辕黄帝, 张海峰, 李光耀."
    ]
    english_translation = translator.translate_cn_to_en(chinese_sentences)
    print(english_translation)

